#ifndef _SAVEASDATASET_H_
#define _SAVEASDATASET_H_

#include "Aimbot.h"

class SaveAsDataset: public Aimbot {
 private:
    std::string save_path;

 public:
    using Armors = std::vector<Armor>;
 public:
    explicit SaveAsDataset(std::string save_path);

    void draw_armors(cv::Mat &img, SaveAsDataset::Armors &armors, int delay);

    int
    label_and_save(int img_cnt, bool back, const std::initializer_list<Armors> &armor_list, bool save_error_only = false);

 private:
    static int correct_label(const cv::Mat &img, int &label, int exist_label, int i);

    static void collect_armors(const std::initializer_list<Armors>& armor_list, Armors& armors);

    int find_exist_label(int img_cnt, int armor_cnt);

    void
    write_img(const cv::Mat &img, int img_cnt, int armor_cnt, int label, int exist_label);

};

#endif //_SAVEASDATASET_H_
